Fix incorrect GPU partitioning and ordering metadata [databricks] by revans2 · Pull Request #16620 · NVIDIA/cudf-spark

revans2 · 2026-10-27T01:35:02Z

GPU projections and aggregations could report partitioning by output position and pass ordering through unchanged. For a query like

spark.range(100).selectExpr('id % 10 as k', 'id as v').groupBy('k').agg(max(col('v')).collect()

The result is a correct 10 rows, but if you persist the result in between the selectExpr and the groupBy like

val data = spark.range(100).selectExpr('id % 10 as k', 'id as v').cache()
data.groupBy('k').agg(max(col('v')).collect()

on the GPU, this causes the ordering and partitioning from the GpuProject to be used and the shuffle before the partial aggregation is not inserted, resulting in 100 rows being returned. This is because the existing code only goes off of position, not inspecting the actual expressions.

This fixes the issue by reusing the installed Spark version’s alias-aware metadata logic for GPU projections and aggregates. It keeps CPU expressions for metadata matching separate from GPU execution expressions, and preserves them through plan copies and projection pruning. If those are not available a conservative fallback happens. This will not be perfect, because the ordering and partitioning can include CPU expressions in them that will not match the GPU expressions, but it will not corrupt the data. In the worst case we may end up with more shuffle operations than before.

Validation:

Standalone Spark 3.5.5 build: 27 Scala and 9 Java tests passed.
Cached aggregate regression passed with AQE on/off under both PCBS and the default cache serializer.
The same metadata implementation passed focused suites on Spark 3.3.4 and 4.1.3.

This appears to have been caused by #14164 and #14264. The first one put in the logic and the second one made it so more operators adopted it. The original logic really only worked for union, and even there I am not 100% sure that it always did.

I will try to run some performance tests, but as this is a data corruption fix it is not strictly required, and if the issue is hit, we will take a performance hit in order to produce the correct result (as a shuffle must be inserted to make the plan correct).

This also adds the internal, test-only spark.rapids.sql.test.forceNonAst setting to force selected expression classes out of AST while retaining GPU expression support. It defaults to an empty list and takes effect only when test mode is enabled.

AI-assisted implementation and testing.

Checklists

Documentation

Updated for new or modified user-facing features or behaviors
No user-facing change

Testing

Added or modified tests to cover new code paths
Covered by existing tests
Not required

Performance

Tests ran and results are added in the PR description
Issue filed with a link in the PR description
Not required

Map partitioning attributes through direct aliases rather than projection output positions. Projecting id % 10 AS k before id AS v otherwise falsely advertised range partitioning on k. Cached scans preserved that claim, so Spark omitted the grouping shuffle and returned per-partition results. Preserve identity and direct CPU/GPU attribute aliases, and let existing reference filtering discard unsupported computed-column partitioning. Keep aggregate and other operator implementations unchanged. Add a compact projection metadata regression and the 100-row cached aggregation reproduction with AQE both enabled and disabled. The new integration regression fails the original jar with CPU 10 / GPU 100 rows. Validation: 25 Scala and 9 Java tests passed on Spark 3.5.5. The new integration regression passed with default cache and PCBS, AQE on/off; full PCBS cache suite: 625 passed, 4 large-data skips. Standalone shell reproductions and controls returned exact results with both serializers. SQL compile/package checks passed on Spark 3.3.4 and 4.0.1/Scala 2.13. Status: READY FOR REVIEW. Verification evidence is retained outside the repository under /data/pcbs-oom/project-partitioning-fix, as requested. Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>

Use the installed Spark version's partitioning and ordering traits instead of maintaining a separate attribute remapping algorithm. Keep original CPU expressions for metadata matching in a second constructor parameter list, following GPU sort and shuffle, so Spark does not treat them as executable expressions. Preserve those expressions across child replacement and projection pruning. Compare metadata directly with native Spark for aliases, computed expressions, ordering prefixes, null ordering and candidate limits. Keep the cached aggregate regression for both AQE settings. Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>

Remove canonicalized-equality assertions that cannot establish anything about second-parameter-list CPU metadata. Keep the direct expression and copy-preservation checks. Collect cached query results before unpersisting in a finally block for each CPU and GPU session, including when query execution fails. Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>

Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>

greptile-apps · 2026-10-27T01:38:58Z

Confidence Score: 5/5

[High impact] No new blocking issue was found in the updated cache tests.

Summary

The PR replaces positional partitioning claims with Spark’s alias-aware partitioning and ordering logic. CPU expressions stay separate from GPU execution expressions and survive plan copies and projection pruning.

GPU projections and aggregates report partitioning and ordering from their expressions.
Projection rewrites keep their metadata in step with the columns.
Tests can keep selected GPU expressions out of AST.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Original CPU expressions] --> C[Spark alias-aware metadata]
  B[GPU expressions] --> D[GPU execution]
  C --> E[Output partitioning and ordering]
  E --> F[Required shuffles and sorts]
  D --> G[Query results]
  F --> G

_{Reviews (8) · Last reviewed commit: "Rework" · Reviewed by Greptile}

Copilot

🟡 Changes recommended

Required performance results or a linked validation issue are missing for changes that can alter shuffle and sort execution.

1 open issue

Add performance validation for Spark metadata change · New

What changed in this PR

Fixes incorrect partitioning and ordering metadata in GPU projections and aggregates, preventing Spark from skipping required shuffles and returning incorrect results.

Changes:

Reuses Spark’s version-specific alias-aware metadata logic while keeping CPU metadata expressions separate from GPU execution expressions.
Preserves metadata through plan copies, projection pruning, and Delta rewrites.
Adds metadata parity tests and a cached-aggregation regression with adaptive execution enabled and disabled.

Shim validation remains INCOMPLETE: the repository snapshot was unavailable, and current-head/base shim-compatibility and shim-related-prs reports could not be verified. The reported focused tests do not establish full affected-profile coverage.

File	Description
tests/src/test/scala/org/apache/spark/sql/rapids/ProjectExprSuite.scala	Tests alias metadata, ordering, candidate limits, and rewrites.
tests/src/test/scala/com/nvidia/spark/rapids/ProjectSplitRetrySuite.scala	Updates projection construction.
tests/src/test/scala/com/nvidia/spark/rapids/HashAggregatesSuite.scala	Updates aggregate construction.
sql-plugin/src/main/spark340/scala/org/apache/spark/sql/rapids/execution/GpuBroadcastNestedLoopJoinExec.scala	Updates join projection construction.
sql-plugin/src/main/spark340/scala/com/nvidia/spark/rapids/shims/GpuAliasAwareExecNodeShims.scala	Adds Spark 3.4+ alias-aware metadata adapters.
sql-plugin/src/main/spark332db/scala/org/apache/spark/sql/rapids/GpuFileFormatWriter.scala	Updates write projection construction.
sql-plugin/src/main/spark330db/scala/org/apache/spark/sql/rapids/execution/GpuBroadcastNestedLoopJoinExec.scala	Updates Databricks join projections.
sql-plugin/src/main/spark330/scala/org/apache/spark/sql/rapids/GpuFileFormatWriter.scala	Updates write projection construction.
sql-plugin/src/main/spark330/scala/org/apache/spark/sql/rapids/execution/GpuBroadcastNestedLoopJoinExec.scala	Updates Spark 3.3 join projections.
sql-plugin/src/main/spark330/scala/com/nvidia/spark/rapids/shims/GpuAliasAwareExecNodeShims.scala	Adds Spark 3.3 alias-aware metadata adapters.
sql-plugin/src/main/scala/org/apache/spark/sql/rapids/execution/GpuRangeBoundaryPlan.scala	Preserves metadata during projection pruning.
sql-plugin/src/main/scala/org/apache/spark/sql/rapids/execution/GpuHashJoin.scala	Updates generated join projections.
sql-plugin/src/main/scala/com/nvidia/spark/rapids/window/GpuWindowExecMeta.scala	Updates generated window projections.
sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAliasAwareOutputExpressions.scala	Converts GPU aliases for metadata matching.
sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAggregateExec.scala	Uses Spark metadata logic with retained CPU expressions.
sql-plugin/src/main/scala/com/nvidia/spark/rapids/basicPhysicalOperators.scala	Replaces positional metadata mapping and preserves projection metadata.
integration_tests/src/main/python/cache_test.py	Adds cached-projection aggregation regression.
delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/DeltaSpark400DB173Provider.scala	Preserves metadata during Databricks Delta rewrites.
delta-lake/delta-spark332db/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuOptimisticTransactionBase.scala	Updates transaction projection construction.
delta-lake/delta-spark330db/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuOptimisticTransactionBase.scala	Updates transaction projection construction.
delta-lake/common/src/main/delta-io/scala/org/apache/spark/sql/delta/rapids/AbstractGpuOptimisticTransactionBase.scala	Updates shared Delta projection construction.
delta-lake/common/src/main/delta-33x-42x/scala/com/nvidia/spark/rapids/delta/common/DeltaProviderBase.scala	Preserves metadata through Delta pruning and merging.
delta-lake/common/src/main/db-350db143-410db183/scala/com/databricks/sql/transaction/tahoe/rapids/GpuOptimisticTransactionBaseCommon.scala	Updates shared Databricks projection construction.

_{Review effort: Balanced}

Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>

revans2 · 2026-10-27T02:05:12Z

build

nartal1

Thanks @revans2 for tracking this down and fixing it. Reusing Spark’s alias-aware logic makes sense. Our earlier tests focused on union behavior and missed this completely.

amahussein

LGTME

A couple of small inline requests.

A test could cover beyond range partitioning.
add a couple of comments to define the contract for the new cpuProjectList argument in the GpuProjectExec.

amahussein · 2026-10-27T22:33:08Z

@revans2 did you check the [AI-AUDIT] GpuUnionExec re-derives partitioning-aware vs concatenating union on every call (SPARK-59122) #16609 ? It sounds like spark-4.1+ GpuUnionExec re-derives its plan at execution.
@nartal1 do you need to update Refactor outputPartitioning for few execs and add tests #14229 striking off item number 2 off the list?

Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>

… contract Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>

Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>

revans2 · 2026-10-28T10:32:37Z

@amahussein your changes should have both been addressed. As for #16690, no I have not done anything with it, but the fix went into 4.2.1, 4.3.X and main. It is mostly a question of do we want to be bug for bug compatible with Apache Spark. Generally I would say no, but because this bug exists only for a single version of Spark and the cost to "fix" it is a 200 line patch in Apache spark, that I don't really know 100% what it would take for us to fix it here means I don't think it is worth doing for 4.2.0. If you disagree we can discuss it in a follow on PR as this fixes data corruption, and I really would like to get it in.

revans2 · 2026-10-28T10:32:54Z

build

revans2 · 2026-10-28T13:48:44Z

The build is failing due to a merge conflict in a test. Working on it now...

Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>

revans2 · 2026-10-28T14:25:32Z

build

amahussein

LGTM

Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>

revans2 · 2026-10-28T19:09:56Z

build

revans2 added 3 commits October 27, 2026 01:18

revans2 requested review from jihoonson and nartal1 and a balanced review from Copilot October 27, 2026 01:35

Copilot started reviewing on behalf of revans2 October 27, 2026 01:35

View session

Merge main into fix-gpu-output-metadata

57bee97

Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>

Copilot AI reviewed Oct 27, 2026

View reviewed changes

Comment thread sql-plugin/src/main/scala/com/nvidia/spark/rapids/basicPhysicalOperators.scala

greptile-apps Bot reviewed Oct 27, 2026

View reviewed changes

Comment thread integration_tests/src/main/python/cache_test.py Outdated

Update copyright years in modified Delta sources

7ef09c6

Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>

greptile-apps Bot reviewed Oct 27, 2026

View reviewed changes

Comment thread sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAggregateExec.scala

Use standard CPU GPU comparison helper for cache regression

1784e2e

Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>

revans2 mentioned this pull request Oct 27, 2026

Reduce Parquet cache heap pressure with smaller row-group records [databricks] #16633

Open

3 of 8 tasks

sameerz added the bug Something isn't working label Oct 27, 2026

amahussein self-requested a review October 27, 2026 16:40

nartal1 previously approved these changes Oct 27, 2026

View reviewed changes

amahussein requested changes Oct 27, 2026

View reviewed changes

Comment thread integration_tests/src/main/python/cache_test.py

Comment thread sql-plugin/src/main/scala/com/nvidia/spark/rapids/basicPhysicalOperators.scala

nartal1 mentioned this pull request Oct 28, 2026

Refactor outputPartitioning for few execs and add tests #14229

Open

revans2 added 3 commits October 28, 2026 02:48

Merge main into fix-gpu-output-metadata

f466cd3

Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>

Expand cached projection metadata regressions and document expression…

2e533fb

… contract Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>

Better fix for test

eeca859

Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>

revans2 dismissed nartal1’s stale review via eeca859 October 28, 2026 10:22

Merge main into fix-gpu-output-metadata

e9b549b

Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>

greptile-apps Bot reviewed Oct 28, 2026

View reviewed changes

Comment thread sql-plugin/src/main/scala/com/nvidia/spark/rapids/RapidsConf.scala

revans2 requested a review from amahussein October 28, 2026 10:40

Upmerge fix

7ba67fe

Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>

amahussein previously approved these changes Oct 28, 2026

View reviewed changes

Comment thread integration_tests/src/main/python/cache_test.py Outdated

Rework

f06c508

Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>

revans2 dismissed amahussein’s stale review via f06c508 October 28, 2026 19:09

amahussein approved these changes Oct 28, 2026

View reviewed changes

amahussein merged commit 6660210 into NVIDIA:main Oct 28, 2026
62 checks passed

Conversation

revans2 commented Oct 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Checklists

Uh oh!

greptile-apps Bot commented Oct 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Confidence Score: 5/5

Summary

Diagram

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Uh oh!

revans2 commented Oct 27, 2026

Uh oh!

nartal1 left a comment

Choose a reason for hiding this comment

Uh oh!

amahussein left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

amahussein commented Oct 27, 2026

Uh oh!

Uh oh!

revans2 commented Oct 28, 2026

Uh oh!

revans2 commented Oct 28, 2026

Uh oh!

revans2 commented Oct 28, 2026

Uh oh!

revans2 commented Oct 28, 2026

Uh oh!

amahussein left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

revans2 commented Oct 28, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

revans2 commented Oct 27, 2026 •

edited

Loading

greptile-apps Bot commented Oct 27, 2026 •

edited

Loading